package net.sourceforge.jsocks;

import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.InterruptedIOException;
import java.io.OutputStream;
import java.io.PrintStream;
import java.io.PushbackInputStream;
import java.net.ConnectException;
import java.net.InetAddress;
import java.net.NoRouteToHostException;
import java.net.ServerSocket;
import java.net.Socket;

import net.sourceforge.jsocks.server.ServerAuthenticator;

/**
 * SOCKS4 and SOCKS5 proxy, handles both protocols simultaniously. Implements
 * all SOCKS commands, including UDP relaying.
 * <p>
 * In order to use it you will need to implement ServerAuthenticator interface.
 * There is an implementation of this interface which does no authentication
 * ServerAuthenticatorNone, but it is very dangerous to use, as it will give
 * access to your local network to anybody in the world. One should never use
 * this authentication scheme unless one have pretty good reason to do so. There
 * is a couple of other authentication schemes in socks.server package.
 * 
 * @see socks.server.ServerAuthenticator
 */
public class ProxyServer implements Runnable {

	ServerAuthenticator auth;
	ProxyMessage msg = null;

	Socket sock = null, remote_sock = null;
	ServerSocket ss = null;
	UDPRelayServer relayServer = null;
	InputStream in, remote_in;
	OutputStream out, remote_out;

	int mode;
	static final int START_MODE = 0;
	static final int ACCEPT_MODE = 1;
	static final int PIPE_MODE = 2;
	static final int ABORT_MODE = 3;

	static final int BUF_SIZE = 8192;

	Thread pipe_thread1, pipe_thread2;
	long lastReadTime;

	protected static int iddleTimeout = 180000; // 3 minutes
	static int acceptTimeout = 180000; // 3 minutes

	static PrintStream log = null;
	static Proxy proxy;

	// Public Constructors
	// ///////////////////

	static final String command_names[] = { "CONNECT", "BIND", "UDP_ASSOCIATE" };

	// Other constructors
	// //////////////////

	static final String command2String(int cmd) {
		if (cmd > 0 && cmd < 4)
			return command_names[cmd - 1];
		else
			return "Unknown Command " + cmd;
	}

	// Public methods
	// ///////////////

	/**
	 * Get proxy.
	 * 
	 * @return Proxy wich is used to handle user requests.
	 */
	public static Proxy getProxy() {
		return proxy;
	}

	static final void log(ProxyMessage msg) {
		log("Request version:" + msg.version + "\tCommand: "
				+ command2String(msg.command));
		log("IP:" + msg.ip + "\tPort:" + msg.port
				+ (msg.version == 4 ? "\tUser:" + msg.user : ""));
	}

	static final void log(String s) {
		if (log != null) {
			log.println(s);
			log.flush();
		}
	}

	/**
	 * Sets the timeout for BIND command, how long the server should wait for
	 * the incoming connection.<br>
	 * Zero timeout implies infinity.<br>
	 * Default timeout is 3 minutes.
	 */
	public static void setAcceptTimeout(int timeout) {
		acceptTimeout = timeout;
	}

	/**
	 * Sets the size of the datagrams used in the UDPRelayServer.<br>
	 * Default size is 64K, a bit more than maximum possible size of the
	 * datagram.
	 */
	public static void setDatagramSize(int size) {
		UDPRelayServer.setDatagramSize(size);
	}

	/**
	 * Sets the timeout for connections, how long shoud server wait for data to
	 * arrive before dropping the connection.<br>
	 * Zero timeout implies infinity.<br>
	 * Default timeout is 3 minutes.
	 */
	public static void setIddleTimeout(int timeout) {
		iddleTimeout = timeout;
	}

	/**
	 * Set the logging stream. Specifying null disables logging.
	 */
	public static void setLog(OutputStream out) {
		if (out == null) {
			log = null;
		} else {
			log = new PrintStream(out, true);
		}

		UDPRelayServer.log = log;
	}

	/**
	 * Set proxy.
	 * <p>
	 * Allows Proxy chaining so that one Proxy server is connected to another
	 * and so on. If proxy supports SOCKSv4, then only some SOCKSv5 requests can
	 * be handled, UDP would not work, however CONNECT and BIND will be
	 * translated.
	 * 
	 * @param p
	 *            Proxy which should be used to handle user requests.
	 */
	public static void setProxy(Proxy p) {
		proxy = p;
		UDPRelayServer.proxy = proxy;
	}

	/**
	 * Sets the timeout for UDPRelay server.<br>
	 * Zero timeout implies infinity.<br>
	 * Default timeout is 3 minutes.
	 */
	public static void setUDPTimeout(int timeout) {
		UDPRelayServer.setTimeout(timeout);
	}

	/**
	 * Creates a proxy server with given Authentication scheme.
	 * 
	 * @param auth
	 *            Authentication scheme to be used.
	 */
	public ProxyServer(ServerAuthenticator auth) {
		this.auth = auth;
	}

	protected ProxyServer(ServerAuthenticator auth, Socket s) {
		this.auth = auth;
		this.sock = s;
		mode = START_MODE;
	}

	private synchronized void abort() {
		if (mode == ABORT_MODE)
			return;
		mode = ABORT_MODE;
		try {
			log("Aborting operation");
			if (remote_sock != null)
				remote_sock.close();
			if (sock != null)
				sock.close();
			if (relayServer != null)
				relayServer.stop();
			if (ss != null)
				ss.close();
			if (pipe_thread1 != null)
				pipe_thread1.interrupt();
			if (pipe_thread2 != null)
				pipe_thread2.interrupt();
		} catch (IOException ioe) {
		}
	}

	private void doAccept() throws IOException {
		Socket s;
		long startTime = System.currentTimeMillis();

		while (true) {
			s = ss.accept();
			if (s.getInetAddress().equals(msg.ip)) {
				// got the connection from the right host
				// Close listenning socket.
				ss.close();
				break;
			} else if (ss instanceof SocksServerSocket) {
				// We can't accept more then one connection
				s.close();
				ss.close();
				throw new SocksException(Proxy.SOCKS_FAILURE);
			} else {
				if (acceptTimeout != 0) { // If timeout is not infinit
					int newTimeout = acceptTimeout
							- (int) (System.currentTimeMillis() - startTime);
					if (newTimeout <= 0)
						throw new InterruptedIOException("In doAccept()");
					ss.setSoTimeout(newTimeout);
				}
				s.close(); // Drop all connections from other hosts
			}
		}

		// Accepted connection
		remote_sock = s;
		remote_in = s.getInputStream();
		remote_out = s.getOutputStream();

		// Set timeout
		remote_sock.setSoTimeout(iddleTimeout);

		log("Accepted from " + s.getInetAddress() + ":" + s.getPort());

		ProxyMessage response;

		if (msg.version == 5)
			response = new Socks5Message(Proxy.SOCKS_SUCCESS,
					s.getInetAddress(), s.getPort());
		else
			response = new Socks4Message(Socks4Message.REPLY_OK,
					s.getInetAddress(), s.getPort());
		response.write(out);
	}

	private void handleException(IOException ioe) {
		// If we couldn't read the request, return;
		if (msg == null)
			return;
		// If have been aborted by other thread
		if (mode == ABORT_MODE)
			return;
		// If the request was successfully completed, but exception happened
		// later
		if (mode == PIPE_MODE)
			return;

		int error_code = Proxy.SOCKS_FAILURE;

		if (ioe instanceof SocksException)
			error_code = ((SocksException) ioe).errCode;
		else if (ioe instanceof NoRouteToHostException)
			error_code = Proxy.SOCKS_HOST_UNREACHABLE;
		else if (ioe instanceof ConnectException)
			error_code = Proxy.SOCKS_CONNECTION_REFUSED;
		else if (ioe instanceof InterruptedIOException)
			error_code = Proxy.SOCKS_TTL_EXPIRE;

		if (error_code > Proxy.SOCKS_ADDR_NOT_SUPPORTED || error_code < 0) {
			error_code = Proxy.SOCKS_FAILURE;
		}

		sendErrorMessage(error_code);
	}

	protected void handleRequest(ProxyMessage msg) throws IOException {
		if (!auth.checkRequest(msg))
			throw new SocksException(Proxy.SOCKS_FAILURE);

		if (msg.ip == null) {
			if (msg instanceof Socks5Message) {
				msg.ip = InetAddress.getByName(msg.host);
			} else
				throw new SocksException(Proxy.SOCKS_FAILURE);
		}
		log(msg);

		switch (msg.command) {
		case Proxy.SOCKS_CMD_CONNECT:
			onConnect(msg);
			break;
		case Proxy.SOCKS_CMD_BIND:
			onBind(msg);
			break;
		case Proxy.SOCKS_CMD_UDP_ASSOCIATE:
			onUDP(msg);
			break;
		default:
			throw new SocksException(Proxy.SOCKS_CMD_NOT_SUPPORTED);
		}
	}

	private void onBind(ProxyMessage msg) throws IOException {
		ProxyMessage response = null;

		if (proxy == null)
			ss = new ServerSocket(0);
		else
			ss = new SocksServerSocket(proxy, msg.ip, msg.port);

		ss.setSoTimeout(acceptTimeout);

		log("Trying accept on " + ss.getInetAddress() + ":" + ss.getLocalPort());

		if (msg.version == 5)
			response = new Socks5Message(Proxy.SOCKS_SUCCESS,
					ss.getInetAddress(), ss.getLocalPort());
		else
			response = new Socks4Message(Socks4Message.REPLY_OK,
					ss.getInetAddress(), ss.getLocalPort());
		response.write(out);

		mode = ACCEPT_MODE;

		pipe_thread1 = Thread.currentThread();
		pipe_thread2 = new Thread(this);
		pipe_thread2.start();

		// Make timeout infinit.
		sock.setSoTimeout(0);
		int eof = 0;

		try {
			while ((eof = in.read()) >= 0) {
				if (mode != ACCEPT_MODE) {
					if (mode != PIPE_MODE)
						return;// Accept failed

					remote_out.write(eof);
					break;
				}
			}
		} catch (EOFException eofe) {
			// System.out.println("EOF exception");
			return;// Connection closed while we were trying to accept.
		} catch (InterruptedIOException iioe) {
			// Accept thread interrupted us.
			// System.out.println("Interrupted");
			if (mode != PIPE_MODE)
				return;// If accept thread was not successfull return.
		} finally {
			// System.out.println("Finnaly!");
		}

		if (eof < 0)// Connection closed while we were trying to accept;
			return;

		// Do not restore timeout, instead timeout is set on the
		// remote socket. It does not make any difference.

		pipe(in, remote_out);
	}

	private void onConnect(ProxyMessage msg) throws IOException {
		Socket s;
		ProxyMessage response = null;

		s = new Socket(msg.ip, msg.port);

		log("Connected to " + s.getInetAddress() + ":" + s.getPort());

		if (msg instanceof Socks5Message) {
			response = new Socks5Message(Proxy.SOCKS_SUCCESS,
					s.getLocalAddress(), s.getLocalPort());
		} else {
			response = new Socks4Message(Socks4Message.REPLY_OK,
					s.getLocalAddress(), s.getLocalPort());

		}
		response.write(out);
		startPipe(s);
	}

	// Private methods
	// ////////////////

	private void onUDP(ProxyMessage msg) throws IOException {
		if (msg.ip.getHostAddress().equals("0.0.0.0"))
			msg.ip = sock.getInetAddress();
		log("Creating UDP relay server for " + msg.ip + ":" + msg.port);
		relayServer = new UDPRelayServer(msg.ip, msg.port,
				Thread.currentThread(), sock, auth);

		ProxyMessage response;

		response = new Socks5Message(Proxy.SOCKS_SUCCESS, relayServer.relayIP,
				relayServer.relayPort);

		response.write(out);

		relayServer.start();

		// Make timeout infinit.
		sock.setSoTimeout(0);
		try {
			while (in.read() >= 0)
				/* do nothing */;
		} catch (EOFException eofe) {
		}
	}

	private void pipe(InputStream in, OutputStream out) throws IOException {
		lastReadTime = System.currentTimeMillis();
		byte[] buf = new byte[BUF_SIZE];
		int len = 0;
		while (len >= 0) {
			try {
				if (len != 0) {
					out.write(buf, 0, len);
					out.flush();
				}
				len = in.read(buf);
				lastReadTime = System.currentTimeMillis();
			} catch (InterruptedIOException iioe) {
				if (iddleTimeout == 0)
					return;// Other thread interrupted us.
				long timeSinceRead = System.currentTimeMillis() - lastReadTime;
				if (timeSinceRead >= iddleTimeout - 1000) // -1s for adjustment.
					return;
				len = 0;

			}
		}
	}

	protected ProxyMessage readMsg(InputStream in) throws IOException {
		PushbackInputStream push_in;
		if (in instanceof PushbackInputStream)
			push_in = (PushbackInputStream) in;
		else
			push_in = new PushbackInputStream(in);

		int version = push_in.read();
		push_in.unread(version);

		ProxyMessage msg;

		if (version == 5) {
			msg = new Socks5Message(push_in, false);
		} else if (version == 4) {
			msg = new Socks4Message(push_in, false);
		} else {
			throw new SocksException(Proxy.SOCKS_FAILURE);
		}
		return msg;
	}

	// Runnable interface
	// //////////////////
	public void run() {
		switch (mode) {
		case START_MODE:
			try {
				startSession();
			} catch (IOException ioe) {
				handleException(ioe);
				// ioe.printStackTrace();
			} finally {
				abort();
				if (auth != null)
					auth.endSession();
				log("Main thread(client->remote)stopped.");
			}
			break;
		case ACCEPT_MODE:
			try {
				doAccept();
				mode = PIPE_MODE;
				pipe_thread1.interrupt(); // Tell other thread that connection
											// have
											// been accepted.
				pipe(remote_in, out);
			} catch (IOException ioe) {
				// log("Accept exception:"+ioe);
				handleException(ioe);
			} finally {
				abort();
				log("Accept thread(remote->client) stopped");
			}
			break;
		case PIPE_MODE:
			try {
				pipe(remote_in, out);
			} catch (IOException ioe) {
			} finally {
				abort();
				log("Support thread(remote->client) stopped");
			}
			break;
		case ABORT_MODE:
			break;
		default:
			log("Unexpected MODE " + mode);
		}
	}

	private void sendErrorMessage(int error_code) {
		ProxyMessage err_msg;
		if (msg instanceof Socks4Message)
			err_msg = new Socks4Message(Socks4Message.REPLY_REJECTED);
		else
			err_msg = new Socks5Message(error_code);
		try {
			err_msg.write(out);
		} catch (IOException ioe) {
		}
	}

	/**
	 * Start the Proxy server at given port.<br>
	 * This methods blocks.
	 */
	public void start(int port) {
		start(port, 5, null);
	}

	/**
	 * Create a server with the specified port, listen backlog, and local IP
	 * address to bind to. The localIP argument can be used on a multi-homed
	 * host for a ServerSocket that will only accept connect requests to one of
	 * its addresses. If localIP is null, it will default accepting connections
	 * on any/all local addresses. The port must be between 0 and 65535,
	 * inclusive. <br>
	 * This methods blocks.
	 */
	public void start(int port, int backlog, InetAddress localIP) {
		try {
			ss = new ServerSocket(port, backlog, localIP);
			log("Starting SOCKS Proxy on:"
					+ ss.getInetAddress().getHostAddress() + ":"
					+ ss.getLocalPort());
			while (true) {
				Socket s = ss.accept();
				log("Accepted from:" + s.getInetAddress().getHostName() + ":"
						+ s.getPort());
				ProxyServer ps = new ProxyServer(auth, s);
				(new Thread(ps)).start();
			}
		} catch (IOException ioe) {
			ioe.printStackTrace();
		} finally {
		}
	}

	private void startPipe(Socket s) {
		mode = PIPE_MODE;
		remote_sock = s;
		try {
			remote_in = s.getInputStream();
			remote_out = s.getOutputStream();
			pipe_thread1 = Thread.currentThread();
			pipe_thread2 = new Thread(this);
			pipe_thread2.start();
			pipe(in, remote_out);
		} catch (IOException ioe) {
		}
	}

	// Private methods
	// ///////////////
	private void startSession() throws IOException {
		sock.setSoTimeout(iddleTimeout);

		try {
			auth = auth.startSession(sock);
		} catch (IOException ioe) {
			log("Auth throwed exception:" + ioe);
			auth = null;
			return;
		}

		if (auth == null) { // Authentication failed
			log("Authentication failed");
			return;
		}

		in = auth.getInputStream();
		out = auth.getOutputStream();

		msg = readMsg(in);
		handleRequest(msg);
	}

	/**
	 * Stop server operation.It would be wise to interrupt thread running the
	 * server afterwards.
	 */
	public void stop() {
		try {
			if (ss != null)
				ss.close();
		} catch (IOException ioe) {
		}
	}
}
